import pandas as pd
from pandas import DataFrame
import DataMiningtools as DMTS

class NaiveBayes (): #封装了朴素贝叶斯算法
    DMT = DMTS.DataMiningTools()

    def __init__(self,train,categorization) -> None:
        self.train = train
        self.categorization = categorization
        self.a_num_arr = self.DMT.AttributeMax(train)
        self.c_num = self.a_num_arr[self.categorization]
        self.CP = self.__CProbability() #存放先验概率
        self.AP = [] #存放条件概率
        col = train.shape[1]

        for i in range (0,col):
            print ('c',i)

            a_num = self.a_num_arr[i]
            self.AP.append([]) #AP[i]为第i个属性
            for a_value in range (0,a_num):
                self.AP[i].append([]) #AP[i][j]为第i个属性的第j个取值
                for c_value in range (0,self.c_num):

                    p = self.__AProbability(i,a_value,c_value,a_num)
                    self.AP[i][a_value].append(p) 
                    #AP[i][j][k] = P(A=Aij|C=Ck)


    def __CProbability(self): #计算先验概率P(Ci)
        train_size = self.train.shape[0]
        CP = [] #CP用来存放每个类的先验概率

        for i in range (0,self.c_num): #初始化CP
            CP.append(0)

        for i in range (0,train_size):

            c = self.train.iloc[i,self.categorization]
            CP[c] += 1

        for i in range (0,self.c_num):

            CP[i] += 1
            CP[i] /= (train_size+self.c_num) #拉普拉斯修正
            
        return CP
                

    def __AProbability (self,attribute,a_value,c_value,a_num): #计算P(A=a|C=c)
        train_size = self.train.shape[0]
        c = 0
        a = 0

        for i in range (0,train_size):
            if self.train.iloc[i,self.categorization] == c_value:
                c += 1
                if self.train.iloc[i,attribute] == a_value:
                    a += 1
        
        a += 1
        c += a_num #拉普拉斯修正
        return a/c


    def XProbability (self,sample,c_value): 
    #计算条件概率P(X|Ci)
        col = sample.shape[1]
        probability = 1

        for i in range (0,col):
            if i == self.categorization: #跳过目标变量
                continue
            
            p = self.AP[i][sample.iloc[0,i]][c_value]

            probability *= p
        
        return probability


    def NaiveBayes_Model(self,sample):
        P_XC = [] #存放样本属于每一类的朴素贝叶斯概率

        for i in range (0,self.c_num):
            CP = self.CP
            p = self.XProbability(sample,i) * CP[i]
            P_XC.append(p)

        max_c = 0
        for i in range (0,self.c_num):
            if P_XC[max_c] < P_XC[i]:
                max_c = i
        return max_c